import heapq
from collections import defaultdict
from itertools import product
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.animation as animation

class Node:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.g = float('inf')  # 到达该节点的成本
        self.rhs = float('inf')  # 启发式成本
        self.parent = None
    
    def __lt__(self, other):
        return (self.g, self.rhs) < (other.g, other.rhs)
    
    def __hash__(self):
        return hash((self.x, self.y))

    def __repr__(self):
        return f"Node({self.x}, {self.y})"

class PriorityQueue:
    def __init__(self):
        self.queue = []
        self.entry_finder = {}

    def push(self, node, key):
        if node in self.entry_finder:
            self.remove(node)
        entry = [key, node]
        self.entry_finder[node] = entry
        heapq.heappush(self.queue, entry)

    def remove(self, node):
        if node in self.entry_finder:
            entry = self.entry_finder.pop(node)

    def pop(self):
        while self.queue:
            priority, node = heapq.heappop(self.queue)
            #不在删除列表里 则可以弹出
            if node  in self.entry_finder:
                del self.entry_finder[node]
                return node
        return None

    def contains(self, node):
        return node in self.entry_finder 

class DStarLite:
    def __init__(self, grid, start, goal,obstacles):
        self.grid = grid
        self.start = start
        self.goal = goal
        self.obstacles = obstacles
        self.km = 0  # 估算的最短路径长度
        self.nodes = {(x, y): Node(x, y) for x in range(len(grid)) for y in range(len(grid[0]))}
        self.open_list = PriorityQueue()
        self.initialize()

    def initialize(self):
        # 初始化目标节点的 rhs 为 0
        self.nodes[self.goal].rhs = 0
        self.update_vertex(self.goal)

    def heuristic(self, u):
        # 使用曼哈顿距离作为启发函数
        return abs(u.x - self.goal[0]) + abs(u.y - self.goal[1])

    def calculate_key(self, u):
        k1 = min(u.g, u.rhs) + self.heuristic(u) + self.km
        k2 = min(u.g, u.rhs)
        return (k1, k2)

    def get_successors(self, u):
        successors = []
        directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]  # 四个方向：上下左右
        for dx, dy in directions:
            x, y = u.x + dx, u.y + dy
            if 0 <= x < len(self.grid) and 0 <= y < len(self.grid[0]) and self.grid[x][y] == 0:
                successors.append((x, y))
        return successors

    def get_predecessors(self, u):
        predecessors = []
        directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]  # 四个方向：上下左右
        for dx, dy in directions:
            x, y = u.x + dx, u.y + dy
            if 0 <= x < len(self.grid) and 0 <= y < len(self.grid[0]) and self.grid[x][y] == 0:
                predecessors.append((x, y))
        return predecessors

    def calculate_cost(self, u, v):
        return 1 if (self.grid[u.x][u.y] == 0 and self.grid[v.x][v.y] == 0) else float('inf')

    # u都是元组
    def update_vertex(self, u):
        if u == self.goal:
            self.nodes[u].rhs = 0
        else:
            successors_costs = [self.calculate_cost(self.nodes[u], self.nodes[v]) + self.nodes[v].g for v in self.get_successors(self.nodes[u])]
            self.nodes[u].rhs = min(successors_costs, default=float('inf'))

        if self.open_list.contains(u):
            self.open_list.remove(u)

        if self.nodes[u].g != self.nodes[u].rhs:
            self.open_list.push(u, self.calculate_key(self.nodes[u]))
        elif self.nodes[u].g > self.nodes[u].rhs:
            u.g = self.nodes[u].rhs
            for s in self.get_predecessors(u):
                self.update_vertex(s)

    def compute_shortest_path(self):
        while self.open_list.queue and (
                self.calculate_key(self.nodes[self.start]) > self.open_list.queue[0][0] or
                self.nodes[self.start].g > self.nodes[self.start].rhs):

            u = self.open_list.pop()
            if u is None:
                return
            u = self.nodes[u]
            k_old = self.calculate_key(u)

            if u.g > u.rhs:
                u.g = u.rhs
                for s in self.get_predecessors(u):
                    self.update_vertex(s)
            else:
                u.g = float('inf')
                self.update_vertex((u.x, u.y))
                for s in self.get_predecessors(u) + [(u.x, u.y)]:
                    print(s)
                    self.update_vertex(s)

    def move_to_goal(self):
        path = []
        current = self.nodes[self.start]
        while tuple([current.x, current.y]) != self.goal:
            path.append((current.x, current.y))
            neighbors = self.get_successors(current)
            next_node = min(neighbors, key=lambda n: self.nodes[n].g, default=None)
            if next_node is None or self.nodes[next_node].g == float('inf'):
                break
            current = self.nodes[next_node]
        path.append((current.x, current.y))
        return path
    def add_obstacle(self, obstacle):
        """
        添加新的障碍物，并更新受影响的节点。
        :param obstacle: 新的障碍物坐标 (x, y)
        """
        x, y = obstacle
        if self.grid[x][y] == 0:  # 如果该位置之前不是障碍物
            self.grid[x][y] = 1  # 将其设置为障碍物
            # self.obstacles.add((x, y))  # 更新障碍物集合

            # 更新与新障碍物相邻的所有节点
            # 定义四个方向：上、右、下、左
            directions = [(-1, 0), (0, 1), (1, 0), (0, -1)]

            for dx, dy in directions:
                nx, ny = x + dx, y + dy
                if 0 <= nx < len(self.grid) and 0 <= ny < len(self.grid[0]):
                    self.update_vertex((nx, ny))

            # 计算新的最短路径
            self.compute_shortest_path()



# Node, PriorityQueue, DStarLite 类保持不变，省略...

def on_click(event, d_star_lite, ax, obstacles, paths):
    if event.xdata is not None and event.ydata is not None:
        x, y = int(event.xdata), int(event.ydata)
        if 0 <= x < len(d_star_lite.grid[0]) and 0 <= y < len(d_star_lite.grid):
            # 添加障碍物
            d_star_lite.add_obstacle((len(d_star_lite.grid) - y - 1, x))
            obstacles.add((len(d_star_lite.grid) - y - 1, x))

            # 重新计算路径
            d_star_lite.compute_shortest_path()
            updated_path = d_star_lite.move_to_goal()

            # 更新路径列表
            paths[0] = updated_path

            # 输出路径
            print("Paths found:")
            for i, path in enumerate(paths):
                print(f"Agent {i+1} Path: {path}")

def visualize_path(paths, map, obstacles, ax=None):
    if ax is None:
        fig, ax = plt.subplots()

    height = len(map)
    width = len(map[0])

    # 绘制栅格
    for y in range(height + 1):
        ax.axhline(y, color='black', linewidth=1)
    for x in range(width + 1):
        ax.axvline(x, color='black', linewidth=1)

    ax.set_xlim(0, width)
    ax.set_ylim(height, 0)

    ax.invert_yaxis()
    ax.set_xticklabels([])
    ax.set_yticklabels([])

    for obstacle in obstacles:
        rect = plt.Rectangle((obstacle[1], height - obstacle[0] - 1), 1, 1, facecolor='gray')
        ax.add_patch(rect)
    
    colors = ['red']
    for i, path in enumerate(paths):
        if path:
            current_position = path[0]
            circle = plt.Circle((current_position[1] + 0.5, height - current_position[0] - 0.5), 0.3, color=colors[i % len(colors)])
            ax.add_patch(circle)
    
    return ax

def update_agents(frame, agents_paths, d_star_lite_instances, obstacles, ax):
    ax.cla()
    visualize_path(agents_paths, d_star_lite_instances[0].grid, obstacles, ax=ax)
    for i, (path, d_star_lite) in enumerate(zip(agents_paths, d_star_lite_instances)):
        if path and len(path) > 1:
            next_position = path.pop(0)
            if next_position != path[0]:
                d_star_lite.start = path[0]
                d_star_lite.compute_shortest_path()
                agents_paths[i] = d_star_lite.move_to_goal()
                print(f"Agent {i+1} moved to {next_position}")
    print("Paths after move:")
    for i, path in enumerate(agents_paths):
        print(f"Agent {i+1} Path: {path}")

if __name__ == '__main__':
    map = [
        [1, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 1, 1, 0],
    ]

    start_positions = [(3, 0)]
    goal_positions = [(0, 3)]
    num_agents = 1

    obstacles = set()
    for x in range(len(map)):
        for y in range(len(map[0])):
            if map[x][y] == 1:
                obstacles.add((x, y))

    agents_paths = []
    d_star_lite_instances = []
    for i in range(num_agents):
        start = start_positions[i]
        goal = goal_positions[i]
        d_star_lite = DStarLite(map, start, goal, obstacles)
        d_star_lite.compute_shortest_path()
        path = d_star_lite.move_to_goal()
        agents_paths.append(path)
        d_star_lite_instances.append(d_star_lite)

    print("Initial Paths found:")
    for i, path in enumerate(agents_paths):
        print(f"Agent {i+1} Path: {path}")

    fig, ax = plt.subplots()
    cid = fig.canvas.mpl_connect('button_press_event', lambda event: on_click(event, d_star_lite_instances[0], ax, obstacles, agents_paths))

    ani = animation.FuncAnimation(fig, update_agents, fargs=(agents_paths, d_star_lite_instances, obstacles, ax),
                                  interval=1000, repeat=True)

    visualize_path(agents_paths, map, obstacles, ax=ax)

    plt.show()